/**
 * @file user_test_api.cc
 * @brief 用户测试API实现
 * @author Intellif
 * @copyright Intellif
 * @date 2021/9/9
 */

#include <tvm/runtime/registry.h>
#include <tvm/runtime/profiling.h>
#include "file_utils.h"
#include <string.h>
#include <dmlc/logging.h>
#include <dmlc/memory_io.h>

namespace tvm {
namespace runtime {

/**
 * @brief RPC空函数调用，callee端
 */
TVM_REGISTER_GLOBAL("rpc.func.call")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    /* do nothing */
});

/**
 * @brief RPC空函数调用，caller端
 * @param args[0] url
 * @param args[1] port
 * @param args[2] test cnt
 */
TVM_REGISTER_GLOBAL("test.rpc.func.call")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    std::string url = args[0];
    int port = args[1];
    int testCnt = args[2];

    Module m = (*Registry::Get("rpc.Connect"))(url, port, "");
    auto rpcTestFunc = m.GetFunction("rpc.func.call");
    tvm::runtime::profiling::Profiler prof{{{kDLCPU, 0},}, {}};

    prof.Start();
    for (int i = 0; i < testCnt; i++) {
        rpcTestFunc();
    }
    prof.Stop();
    auto rep = prof.Report();
    std::cout << rep.as<tvm::runtime::profiling::ReportNode>()->AsJSON() << std::endl;

    *rv = true;
});

/**
 * @brief 文件传输，caller端
 * @param args[0] url
 * @param args[1] port
 * @param args[2] file path
 */
TVM_REGISTER_GLOBAL("test.rpc.upload.file")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    std::string url = args[0];
    int port = args[1];
    std::string filePath = args[2];

    Module m = (*Registry::Get("rpc.Connect"))(url, port, "");
    auto rpcTestFunc = m.GetFunction("tvm.rpc.server.upload");
    tvm::runtime::profiling::Profiler prof{{{kDLCPU, 0},}, {}};
    std::string strFileData;

    prof.Start();
    LoadBinaryFromFile(filePath, &strFileData);
    TVMByteArray byteArray{strFileData.data(), strFileData.size()};
    rpcTestFunc("/tmp/tempfile", byteArray);
    prof.Stop();
    auto rep = prof.Report();
    std::cout << rep.as<tvm::runtime::profiling::ReportNode>()->AsJSON() << std::endl;

    *rv = true;
});

/**
 * @brief NDArray传输测试，创建一个NDArray
 * @param args[0] ndarray size
 */
TVM_REGISTER_GLOBAL("rpc.create.ndarray")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    int ndSize = args[0];
    NDArray transfer_test_nd = NDArray::Empty({ndSize}, String2DLDataType("int8"), Device{kDLCPU, 0});
    int* idata = static_cast<int*>(transfer_test_nd->data);
    for (int i = 0; i < ndSize / 4; i++) {
        idata[i] =  i;
    }

    *rv = transfer_test_nd;
});

/**
 * @brief NDArray传输测试，server copy from client
 * @param args[0] url
 * @param args[1] port
 * @param args[2] ndarray size
 * @param args[3] test cnt
 */
TVM_REGISTER_GLOBAL("test.rpc.ndarray.copy")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    std::string url = args[0];
    int port = args[1];
    int ndSize = args[2];
    int testCnt = args[3];

    //获取远端NDArray对象
    Module m = (*Registry::Get("rpc.Connect"))(url, port, "");
    auto rpcTestFunc = m.GetFunction("rpc.create.ndarray");
    NDArray ndServer = rpcTestFunc(ndSize);

    //创建本地NDArray对象
    NDArray ndClient = NDArray::Empty({ndSize}, String2DLDataType("int8"), Device{kDLCPU, 0});

    {
        tvm::runtime::profiling::Profiler prof{{{kDLCPU, 0},}, {}};
        prof.Start();
        for (int i = 0; i < testCnt; i++) {
            ndServer.CopyFrom(ndClient);
        }
        prof.Stop();
        auto rep = prof.Report();
        std::cout << rep.as<tvm::runtime::profiling::ReportNode>()->AsJSON() << std::endl;
    }
    

    {
        tvm::runtime::profiling::Profiler prof{{{kDLCPU, 0},}, {}};
        prof.Start();
        for (int i = 0; i < testCnt; i++) {
            ndClient.CopyFrom(ndServer);
        }
        prof.Stop();
        auto rep = prof.Report();
        std::cout << rep.as<tvm::runtime::profiling::ReportNode>()->AsJSON() << std::endl;
    }
    
    *rv = true;
});

/**
 * @brief 写入NDArray到文件中
 * @param args[0] file path
 * @param args[1] NDArray object
 */
TVM_REGISTER_GLOBAL("write.ndarray.to.file")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    std::string filePath = args[0];

    std::string strData;
    ::dmlc::MemoryStringStream strmIn(&strData);

    // 通过RPC传递NDArray，会自动转为DLTensorHandle。本地调用传递NDArray，则保持不变
    // 但不管args是NDArray还是DLTensorHandle，都可以统一用DLTensor获取其值
    DLTensor* tensor = args[1];
    SaveDLTensor(&strmIn, tensor);
    SaveBinaryToFile(filePath, strData);

    *rv = true;
});

/**
 * @brief 从文件中读取NDArray
 * @param args[0] file path
 * @return 从文件中恢复出的NDArray object
 */
TVM_REGISTER_GLOBAL("read.ndarray.from.file")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    std::string filePath = args[0];
    
    std::string strData;
    LoadBinaryFromFile(filePath, &strData);

    ::dmlc::MemoryStringStream strmOut(&strData);
    NDArray nd;
    nd.Load(&strmOut);

    *rv = nd;
});

/**
 * @brief 从string读取NDArray
 * @param args[0] NDArray data string
 * @return 从string恢复出的NDArray object
 */
TVM_REGISTER_GLOBAL("read.ndarray.from.string")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    std::string strData = args[0];
    ::dmlc::MemoryStringStream strmOut(&strData);
    NDArray nd;
    nd.Load(&strmOut);

    *rv = nd;
});

}
}